{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b2294610-9701-4591-9ac8-b621cca55247",
   "metadata": {},
   "source": [
    "# 图像弹性增广\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80b904e6-cc6f-4420-ae79-57a1310a73d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding:utf-8 -*-\n",
    "import cv2\n",
    "import numpy as np\n",
    "from scipy.ndimage import gaussian_filter\n",
    "from scipy.ndimage import map_coordinates\n",
    "\n",
    "\n",
    "def elastic_transform(image, alpha, sigma,\n",
    "                      alpha_affine, random_state=None):\n",
    "\n",
    "    if random_state is None:\n",
    "        random_state = np.random.RandomState(None)\n",
    "\n",
    "    shape = image.shape\n",
    "    shape_size = shape[:2]\n",
    "    # Random affine\n",
    "    center_square = np.float32(shape_size) // 2\n",
    "    square_size = min(shape_size) // 3\n",
    "    # pts1: 仿射变换前的点(3个点)\n",
    "    pts1 = np.float32([center_square + square_size,\n",
    "                       [center_square[0] + square_size,\n",
    "                        center_square[1] - square_size],\n",
    "                       center_square - square_size])\n",
    "    # pts2: 仿射变换后的点\n",
    "    pts2 = pts1 + random_state.uniform(-alpha_affine, alpha_affine,\n",
    "                                       size=pts1.shape).astype(np.float32)\n",
    "    # 仿射变换矩阵\n",
    "    M = cv2.getAffineTransform(pts1, pts2)\n",
    "    # 对image进行仿射变换.\n",
    "    imageB = cv2.warpAffine(image, M, shape_size[::-1], borderMode=cv2.BORDER_REFLECT_101)\n",
    "\n",
    "    # generate random displacement fields\n",
    "    # random_state.rand(*shape)会产生一个和shape一样打的服从[0,1]均匀分布的矩阵\n",
    "    # *2-1是为了将分布平移到[-1, 1]的区间, alpha是控制变形强度的变形因子\n",
    "    dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha\n",
    "    dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha\n",
    "    # generate meshgrid\n",
    "    x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))\n",
    "    # x+dx,y+dy\n",
    "    indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx, (-1, 1))\n",
    "    # bilinear interpolation\n",
    "    imageC = map_coordinates(imageB, indices, order=1, mode='constant').reshape(shape)\n",
    "\n",
    "    return imageC\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    img_path = '/storage/data/images/cat-01.jpg'\n",
    "    imageA = cv2.imread(img_path)\n",
    "    img_show = imageA.copy()\n",
    "    imageA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)\n",
    "    # Apply elastic transform on image\n",
    "    imageC = elastic_transform(imageA, imageA.shape[1] * 2,\n",
    "                                   imageA.shape[1] * 0.08,\n",
    "                                   imageA.shape[1] * 0.08)\n",
    "\n",
    "    cv2.namedWindow(\"img_a\", 0)\n",
    "    cv2.imshow(\"img_a\", img_show)\n",
    "    cv2.namedWindow(\"img_c\", 0)\n",
    "    cv2.imshow(\"img_c\", imageC)\n",
    "    # cv2.waitKey(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bedc386e-78e6-4313-86c2-43939f5de596",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py3.12_torch22",
   "language": "python",
   "name": "py3.12_torch22"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
